import asyncio

import aiohttp
import redis
import json
import uuid
import time
from typing import Dict, Optional, List
from pydantic import BaseModel
from nacos import NacosClient
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
from aiobreaker import CircuitBreaker, CircuitBreakerError, CircuitBreakerState as State
from prometheus_client import Counter, Histogram, Gauge, start_http_server
from loguru import logger

# -------------------------- 1. 监控指标定义（无修改） --------------------------
FLOW_TOTAL = Counter(
    "distributed_flow_total",
    "Total number of distributed agent flows",
    ["flow_type"]
)
FLOW_SUCCESS = Counter(
    "distributed_flow_success",
    "Number of successful distributed agent flows",
    ["flow_type"]
)
FLOW_FAILED = Counter(
    "distributed_flow_failed",
    "Number of failed distributed agent flows",
    ["flow_type", "error_type"]
)
FLOW_DURATION = Histogram(
    "distributed_flow_duration_seconds",
    "Duration of distributed agent flows",
    ["flow_type"]
)
LOCK_COMPETITION = Counter(
    "distributed_lock_competition",
    "Number of distributed lock competition events",
    ["resource"]
)
AGENT_HEALTH = Gauge(
    "agent_service_health",
    "Health status of agent services (1=healthy, 0=unhealthy)",
    ["agent_name"]
)


# -------------------------- 2. 结构化状态模型（修复反斜杠问题） --------------------------
class FlowState(BaseModel):
    flow_id: str  # 流程ID
    trace_id: str  # 链路ID（串联跨节点日志）
    user_query: str  # 用户查询
    selected_agents: List[str] = []  # 选中的Agent列表
    agent_results: Dict[str, str] = {}  # Agent调用结果
    created_at: float = time.time()  # 创建时间
    updated_at: float = time.time()  # 更新时间
    final_result: str = ""  # 最终整合结果

    def dict(self, *args, **kwargs):
        # 重写dict方法，确保时间戳序列化（无修改）
        data = super().dict(*args, **kwargs)
        data["created_at"] = float(self.created_at)
        data["updated_at"] = float(self.updated_at)
        return data

    @classmethod
    def from_dict(cls, data: Dict):
        # 修复：用int()替代反斜杠，Python 3.11 f-string不允许反斜杠
        # 原错误代码：created_at=data.get("created_at", time.time()\n),
        return cls(
            flow_id=data["flow_id"],
            trace_id=data["trace_id"],
            user_query=data["user_query"],
            selected_agents=data.get("selected_agents", []),
            agent_results=data.get("agent_results", {}),
            # 修复点：直接用time.time()，去掉多余反斜杠（反斜杠是笔误）
            created_at=data.get("created_at", time.time()),
            updated_at=data.get("updated_at", time.time()),
            final_result=data.get("final_result", "")
        )


# -------------------------- 3. Redis分布式状态（无修改） --------------------------
class RedisDistributedState:
    def __init__(self, redis_host: str = "127.0.0.1", redis_port: int = 6379):
        self.redis = redis.Redis(host=redis_host, port=redis_port, decode_responses=True)
        self.prefix = "distributed_flow:"
        self.lock_prefix = "lock:"
        self.state_ttl = 3600  # 状态1小时过期
        self.lock_renew_interval = 3  # 锁续约间隔（秒）

    def generate_flow_id(self) -> str:
        return f"flow_{uuid.uuid4().hex[:16]}"

    def generate_trace_id(self) -> str:
        return f"trace_{uuid.uuid4().hex[:24]}"

    def save_state(self, flow_id: str, state: FlowState) -> bool:
        try:
            state_key = f"{self.prefix}{flow_id}"
            state.updated_at = time.time()
            self.redis.set(state_key, json.dumps(state.dict()))
            self.redis.expire(state_key, self.state_ttl)
            logger.debug(f"[Redis] 保存状态成功，flow_id={flow_id}")
            return True
        except Exception as e:
            logger.error(f"[Redis] 保存状态失败，flow_id={flow_id}，err={str(e)}")
            return False

    def get_state(self, flow_id: str) -> Optional[FlowState]:
        try:
            state_key = f"{self.prefix}{flow_id}"
            state_str = self.redis.get(state_key)
            if not state_str:
                logger.warning(f"[Redis] 状态不存在，flow_id={flow_id}")
                return None
            return FlowState.from_dict(json.loads(state_str))
        except Exception as e:
            logger.error(f"[Redis] 获取状态失败，flow_id={flow_id}，err={str(e)}")
            return None

    def get_lock(self, flow_id: str, timeout: int = 10) -> Optional["RedisLock"]:
        lock_key = f"{self.lock_prefix}{self.prefix}{flow_id}"
        lock_value = str(uuid.uuid4())
        if self.redis.set(lock_key, lock_value, nx=True, px=timeout * 1000):
            logger.debug(f"[Redis锁] 获取成功，flow_id={flow_id}，timeout={timeout}s")
            return RedisLock(
                redis_client=self.redis,
                lock_key=lock_key,
                lock_value=lock_value,
                renew_interval=self.lock_renew_interval,
                timeout=timeout
            )
        else:
            LOCK_COMPETITION.labels(resource=f"flow_{flow_id}").inc()
            logger.warning(f"[Redis锁] 获取失败（竞争），flow_id={flow_id}")
            return None


class RedisLock:
    def __init__(self, redis_client, lock_key: str, lock_value: str, renew_interval: int, timeout: int):
        self.redis = redis_client
        self.lock_key = lock_key
        self.lock_value = lock_value
        self.renew_interval = renew_interval
        self.timeout = timeout
        self.renew_task = None
        self.is_released = False

    async def _renew_lock(self):
        while not self.is_released:
            try:
                renew_script = """
                if redis.call('get', KEYS[1]) == ARGV[1] then
                    return redis.call('pexpire', KEYS[1], ARGV[2])
                else
                    return 0
                end
                """
                result = self.redis.eval(renew_script, 1, self.lock_key, self.lock_value, self.timeout * 1000)
                if result == 0:
                    logger.warning(f"[Redis锁续约] 锁已丢失，key={self.lock_key}")
                    break
                logger.debug(f"[Redis锁续约] 成功，key={self.lock_key}")
                await asyncio.sleep(self.renew_interval)
            except Exception as e:
                logger.error(f"[Redis锁续约] 失败，key={self.lock_key}，err={str(e)}")
                await asyncio.sleep(1)

    def acquire(self):
        self.renew_task = asyncio.create_task(self._renew_lock())

    def release(self):
        if self.is_released:
            return
        try:
            self.is_released = True
            if self.renew_task:
                self.renew_task.cancel()
            unlock_script = """
            if redis.call('get', KEYS[1]) == ARGV[1] then
                return redis.call('del', KEYS[1])
            else
                return 0
            end
            """
            result = self.redis.eval(unlock_script, 1, self.lock_key, self.lock_value)
            if result == 1:
                logger.debug(f"[Redis锁] 释放成功，key={self.lock_key}")
            else:
                logger.warning(f"[Redis锁] 释放失败（锁已丢失），key={self.lock_key}")
        except Exception as e:
            logger.error(f"[Redis锁] 释放异常，key={self.lock_key}，err={str(e)}")

    def __del__(self):
        self.release()


# -------------------------- 4. Nacos服务发现（无修改） --------------------------
class NacosServiceDiscovery:
    # 原代码：server_addresses="127.0.0.1:8848"
    # 修改后：使用你的Nacos IP
    def __init__(self, server_addresses: str = "192.168.150.101:884", namespace: str = "public"):
        self.nacos = NacosClient(server_addresses=server_addresses, namespace=namespace)
        self.instance_cache = {}  # 实例缓存（减少Nacos调用）
        self.cache_ttl = 30  # 缓存30秒过期

    def get_healthy_instances(self, service_name: str) -> List[Dict]:
        try:
            cache_key = service_name
            now = time.time()
            if cache_key in self.instance_cache:
                cache_time, instances = self.instance_cache[cache_key]
                if now - cache_time < self.cache_ttl:
                    healthy_instances = [ins for ins in instances if ins["healthy"]]
                    AGENT_HEALTH.labels(agent_name=service_name).set(1 if healthy_instances else 0)
                    return healthy_instances

            instances = self.nacos.list_naming_instance(service_name)["hosts"]
            self.instance_cache[cache_key] = (now, instances)
            healthy_instances = [ins for ins in instances if ins["healthy"]]
            AGENT_HEALTH.labels(agent_name=service_name).set(1 if healthy_instances else 0)
            logger.debug(f"[Nacos] 拉取实例成功，service={service_name}，健康实例数={len(healthy_instances)}")
            return healthy_instances
        except Exception as e:
            AGENT_HEALTH.labels(agent_name=service_name).set(0)
            logger.error(f"[Nacos] 拉取实例失败，service={service_name}，err={str(e)}")
            return []

    def choose_instance(self, service_name: str) -> Optional[Dict]:
        instances = self.get_healthy_instances(service_name)
        if not instances:
            return None
        hash_key = f"{service_name}_{time.time() // 10}"
        index = hash(hash_key) % len(instances)
        return instances[index]


# -------------------------- 5. Agent调用器（无修改） --------------------------
class AgentInvoker:
    def __init__(self, http_client, nacos_discovery: NacosServiceDiscovery):
        self.http_client = http_client
        self.nacos = nacos_discovery
        self.breaker_map = {
            "rag_agent_service": self._create_circuit_breaker("rag_agent_service"),
            "mcp_agent_service": self._create_circuit_breaker("mcp_agent_service")
        }

    def _create_circuit_breaker(self, agent_name: str) -> CircuitBreaker:
        def on_failure(failure: Exception):
            logger.error(f"[熔断器] Agent调用失败触发熔断，agent={agent_name}，err={str(failure)}")

        def on_state_change(state: State):
            logger.warning(f"[熔断器] Agent状态变更，agent={agent_name}，state={state.name}")
            AGENT_HEALTH.labels(agent_name=f"{agent_name}_circuit").set(1 if state == State.OPEN else 0)

        return CircuitBreaker(
            fail_max=5,
            fail_timeout=5,
            on_failure=on_failure,
            on_state_change=on_state_change
        )

    @retry(
        stop=stop_after_attempt(2),
        wait=wait_exponential(multiplier=1, min=1, max=3),
        retry=retry_if_exception_type((asyncio.TimeoutError, aiohttp.ClientError)),
        reraise=True
    )
    async def _call_agent_http(self, agent_url: str, flow_id: str, trace_id: str, query: str) -> str:
        try:
            async with self.http_client.post(
                    agent_url,
                    json={"flow_id": flow_id, "query": query, "trace_id": trace_id},
                    timeout=10
            ) as resp:
                resp_data = await resp.json()
                if resp.status != 200:
                    raise Exception(f"HTTP状态异常：{resp.status}，msg={resp_data.get('message', '')}")
                return resp_data.get("data", "无返回结果")
        except Exception as e:
            logger.error(f"[Agent HTTP调用] 失败，url={agent_url}，err={str(e)}")
            raise

    async def invoke_agent(self, agent_name: str, flow_id: str, trace_id: str, query: str) -> str:
        breaker = self.breaker_map.get(agent_name, None)
        if not breaker:
            return f"不支持的Agent：{agent_name}"

        try:
            with breaker:
                instance = self.nacos.choose_instance(agent_name)
                if not instance:
                    raise Exception("无健康实例可用")

                agent_type = agent_name.split("_")[0]
                agent_url = f"http://{instance['ip']}:{instance['port']}/agent/{agent_type}/invoke"
                result = await self._call_agent_http(agent_url, flow_id, trace_id, query)
                logger.info(f"[Agent调用] 成功，agent={agent_name}，flow_id={flow_id}")
                return result
        except CircuitBreakerError:
            return f"Agent[{agent_name}]已熔断，暂时无法调用"
        except Exception as e:
            return f"Agent[{agent_name}]调用失败：{str(e)}"